﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.CompilerServices;
using System.Threading;

namespace Util
{
    public class Cache<K, V> where K:class where V:class
    {
        private readonly ConditionalWeakTable<K, V> cache;
        private readonly Func<K, V> computation;
        private readonly Dictionary<K, HashSet<Action<V>>> computing;

        public Cache(Func<K, V> computation)
        {
            this.cache = new ConditionalWeakTable<K, V>();
            this.computation = computation;
            this.computing = new Dictionary<K, HashSet<Action<V>>>();
        }

        public V this[K input]
        {
            get {
                V res;
                bool needToCompute;
                object localLock = null;
                bool waitingForRes = true;
                lock (computing)
                {
                    if (cache.TryGetValue(input, out res))
                    {
                        return res;
                    }
                    else
                    {
                        HashSet<Action<V>> actions;
                        if (computing.TryGetValue(input, out actions))
                        {
                            needToCompute = false;
                            localLock = new object();
                            Action<V> action = ((computedVal) =>
                            {
                                res = computedVal;
                                lock (localLock)
                                {
                                    waitingForRes = false;
                                    Monitor.PulseAll(localLock);
                                }
                            });
                            actions.Add(action);
                        }
                        else
                        {
                            needToCompute = true;
                            computing[input] = new HashSet<Action<V>>();
                        }
                    }
                }
                if (needToCompute)
                {
                    res = computation(input);
                    cache.Add(input, res);

                    lock (computing)
                    {
                        foreach (var remoteAction in computing[input])
                        {
                            remoteAction(res);
                        }
                    }
                }
                else
                {
                    lock (localLock)
                    {
                        if (waitingForRes)
                        {
                            Monitor.Wait(localLock);
                        }
                    }
                }
                return res;
            }
        }
    }
}
